package com.sprint.questai.model.extraModel.embedding;

import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.DocumentSplitter;
import dev.langchain4j.data.document.DocumentTransformer;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.data.segment.TextSegmentTransformer;
import dev.langchain4j.internal.Utils;
import dev.langchain4j.internal.ValidationUtils;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.spi.ServiceHelper;
import dev.langchain4j.spi.data.document.splitter.DocumentSplitterFactory;
import dev.langchain4j.spi.model.embedding.EmbeddingModelFactory;
import dev.langchain4j.store.embedding.EmbeddingStore;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.*;
import java.util.stream.Collectors;

/**
 * @Author: mayixiang
 * @Date: 2024/6/22 22:23
 * @Description:
 */
public class NewEmbeddingStoreIngestor {
    private static final Logger log = LoggerFactory.getLogger(NewEmbeddingStoreIngestor.class);
    private final DocumentTransformer documentTransformer;
    private final DocumentSplitter documentSplitter;
    private final TextSegmentTransformer textSegmentTransformer;
    private final EmbeddingModel embeddingModel;
    private final EmbeddingStore<TextSegment> embeddingStore;

    public NewEmbeddingStoreIngestor(DocumentTransformer documentTransformer, DocumentSplitter documentSplitter, TextSegmentTransformer textSegmentTransformer, EmbeddingModel embeddingModel, EmbeddingStore<TextSegment> embeddingStore) {
        this.documentTransformer = documentTransformer;
        this.documentSplitter = (DocumentSplitter) Utils.getOrDefault(documentSplitter, NewEmbeddingStoreIngestor::loadDocumentSplitter);
        this.textSegmentTransformer = textSegmentTransformer;
        this.embeddingModel = (EmbeddingModel) ValidationUtils.ensureNotNull((EmbeddingModel)Utils.getOrDefault(embeddingModel, NewEmbeddingStoreIngestor::loadEmbeddingModel), "embeddingModel");
        this.embeddingStore = (EmbeddingStore)ValidationUtils.ensureNotNull(embeddingStore, "embeddingStore");
    }

    private static DocumentSplitter loadDocumentSplitter() {
        Collection<DocumentSplitterFactory> factories = ServiceHelper.loadFactories(DocumentSplitterFactory.class);
        if (factories.size() > 1) {
            throw new RuntimeException("Conflict: multiple document splitters have been found in the classpath. Please explicitly specify the one you wish to use.");
        } else {
            Iterator var1 = factories.iterator();
            if (var1.hasNext()) {
                DocumentSplitterFactory factory = (DocumentSplitterFactory)var1.next();
                DocumentSplitter documentSplitter = factory.create();
                log.debug("Loaded the following document splitter through SPI: {}", documentSplitter);
                return documentSplitter;
            } else {
                return null;
            }
        }
    }

    private static EmbeddingModel loadEmbeddingModel() {
        Collection<EmbeddingModelFactory> factories = ServiceHelper.loadFactories(EmbeddingModelFactory.class);
        if (factories.size() > 1) {
            throw new RuntimeException("Conflict: multiple embedding models have been found in the classpath. Please explicitly specify the one you wish to use.");
        } else {
            Iterator var1 = factories.iterator();
            if (var1.hasNext()) {
                EmbeddingModelFactory factory = (EmbeddingModelFactory)var1.next();
                EmbeddingModel embeddingModel = factory.create();
                log.debug("Loaded the following embedding model through SPI: {}", embeddingModel);
                return embeddingModel;
            } else {
                return null;
            }
        }
    }

    public static void ingest(Document document, EmbeddingStore<TextSegment> embeddingStore) {
        builder().embeddingStore(embeddingStore).build().ingest(document);
    }

    public static void ingest(List<Document> documents, EmbeddingStore<TextSegment> embeddingStore) {
        builder().embeddingStore(embeddingStore).build().ingest(documents);
    }

    public void ingest(Document document) {
        this.ingest(Collections.singletonList(document));
    }

    public void ingest(Document... documents) {
        this.ingest(Arrays.asList(documents));
    }
    public String ingestR(Document document) {
        return this.ingestR(Collections.singletonList(document)).get(0);
    }

    public List<String> ingestR(List<Document> documents) {
        log.debug("Starting to ingest {} documents", documents.size());
        if (this.documentTransformer != null) {
            documents = this.documentTransformer.transformAll(documents);
            log.debug("Documents were transformed into {} documents", documents.size());
        }

        List segments;
        if (this.documentSplitter != null) {
            segments = this.documentSplitter.splitAll(documents);
            log.debug("Documents were split into {} text segments", segments.size());
        } else {
            segments = (List)documents.stream().map(Document::toTextSegment).collect(Collectors.toList());
        }

        if (this.textSegmentTransformer != null) {
            segments = this.textSegmentTransformer.transformAll(segments);
            log.debug("Text segments were transformed into {} text segments", documents.size());
        }

        log.debug("Starting to embed {} text segments", segments.size());
        List<Embedding> embeddings = (List)this.embeddingModel.embedAll(segments).content();
        log.debug("Finished embedding {} text segments", segments.size());
        log.debug("Starting to store {} text segments into the embedding store", segments.size());
        List<String> list = this.embeddingStore.addAll(embeddings, segments);
        log.debug("Finished storing {} text segments into the embedding store", segments.size());
        return list;
    }

    public void ingest(List<Document> documents) {
        log.debug("Starting to ingest {} documents", documents.size());
        if (this.documentTransformer != null) {
            documents = this.documentTransformer.transformAll(documents);
            log.debug("Documents were transformed into {} documents", documents.size());
        }

        List segments;
        if (this.documentSplitter != null) {
            segments = this.documentSplitter.splitAll(documents);
            log.debug("Documents were split into {} text segments", segments.size());
        } else {
            segments = (List)documents.stream().map(Document::toTextSegment).collect(Collectors.toList());
        }

        if (this.textSegmentTransformer != null) {
            segments = this.textSegmentTransformer.transformAll(segments);
            log.debug("Text segments were transformed into {} text segments", documents.size());
        }

        log.debug("Starting to embed {} text segments", segments.size());
        List<Embedding> embeddings = (List)this.embeddingModel.embedAll(segments).content();
        log.debug("Finished embedding {} text segments", segments.size());
        log.debug("Starting to store {} text segments into the embedding store", segments.size());
        this.embeddingStore.addAll(embeddings, segments);
        log.debug("Finished storing {} text segments into the embedding store", segments.size());
    }

    public static NewEmbeddingStoreIngestor.Builder builder() {
        return new NewEmbeddingStoreIngestor.Builder();
    }

    public static class Builder {
        private DocumentTransformer documentTransformer;
        private DocumentSplitter documentSplitter;
        private TextSegmentTransformer textSegmentTransformer;
        private EmbeddingModel embeddingModel;
        private EmbeddingStore<TextSegment> embeddingStore;

        public Builder() {
        }

        public NewEmbeddingStoreIngestor.Builder documentTransformer(DocumentTransformer documentTransformer) {
            this.documentTransformer = documentTransformer;
            return this;
        }

        public NewEmbeddingStoreIngestor.Builder documentSplitter(DocumentSplitter documentSplitter) {
            this.documentSplitter = documentSplitter;
            return this;
        }

        public NewEmbeddingStoreIngestor.Builder textSegmentTransformer(TextSegmentTransformer textSegmentTransformer) {
            this.textSegmentTransformer = textSegmentTransformer;
            return this;
        }

        public NewEmbeddingStoreIngestor.Builder embeddingModel(EmbeddingModel embeddingModel) {
            this.embeddingModel = embeddingModel;
            return this;
        }

        public NewEmbeddingStoreIngestor.Builder embeddingStore(EmbeddingStore<TextSegment> embeddingStore) {
            this.embeddingStore = embeddingStore;
            return this;
        }

        public NewEmbeddingStoreIngestor build() {
            return new NewEmbeddingStoreIngestor(this.documentTransformer, this.documentSplitter, this.textSegmentTransformer, this.embeddingModel, this.embeddingStore);
        }
    }
}